#!/usr/bin/perl
# processes alanine scanning + dssp data
# by andy
# helix-only version

use strict;
use warnings;

use DBI;
use List::Util qw(sum);

use Data::Dumper;

# Skip DB for now...
#    # Database parameters - change these as appropriate.
#    $username = '';$password = '';$database = '';$hostname = '';
#    $tablename = "chains" # change this too
#    $dbh = DBI->connect("dbi:mysql:database=$database;" .
#       "host=$hostname;port=3306", $username, $password);

# Naccess path
my $naccess = "/scratch/NETID/pdb_interact/Naccess/naccess";

use Class::Struct;
struct(residue => [chain => '$', pdbresnum => '$', aa => '$', rosresnum => '$', ddg => '$', dssp => '$', dsasa => '$']);

my %aa_3_to_1 = ("ala", 'A',
                 "cys", 'C',
                 "asp", 'D',
                 "glu", 'E',
                 "phe", 'F',
                 "gly", 'G',
                 "his", 'H',
                 "ile", 'I',
                 "lys", 'K',
                 "leu", 'L',
                 "met", 'M',
                 "asn", 'N',
                 "pro", 'P',
                 "gln", 'Q',
                 "arg", 'R',
                 "ser", 'S',
                 "thr", 'T',
                 "val", 'V',
                 "trp", 'W',
                 "tyr", 'Y',
                );

my %chains;
my %chain1;
my %chain2;
my @this_chains_keys;
my %hotspots;
my $pdb_code;
my $first_chain;
my $second_chain;

open(LIST, "<interface_chains_pdb.txt")
    or die "ERROR: Could not open interface_chains_pdb.txt\n";

open(SQLOUT, ">010413.sql");

CODELOOP: foreach my $pdb_file (<LIST>) {
    
    chomp $pdb_file;

    #print STDERR "Parsing $pdb_file\n";
    
    # Isolate PDB code from line in interface_chains_pdb.txt
    # Go from "pdb_partner_files/1A0H_A_B.pdb" to "1A0H_A_B"
    $pdb_code = $pdb_file;
    $pdb_code =~ s/pdb_partner_files\///;
    $pdb_code =~ s/\.pdb//;
    
    # Figure out codes for first and second chains.
    if ($pdb_code =~ /.{4}_(.)_(.)/) {
        $first_chain  = $1;
        $second_chain = $2;
    }
    
    ##-----------------------------------------------------------------------##
    ## Reading in residue data from PDB file --------------------------------##
    ##-----------------------------------------------------------------------##
    
    # Hashes for both chains
    undef(%chain1);
    undef(%chain2);
    
    unless (open(PDB, "<$pdb_file")) {
        print STDERR "ERROR: Could not open PDB file $pdb_file.\n";
        next;
    }
    
    my $old_pdb_res = 0;
    foreach (<PDB>) {
        my @pdb_elem = split;
        
        # Skip if not an ATOM line
        if ($pdb_elem[0] ne 'ATOM') {
            next;
        }
    
    my $pdb_res;
    my $pdb_chain;
        if (substr $pdb_elem[4],0,1 eq $first_chain and substr $pdb_elem[4], 1 eq $old_pdb_res+1) {
            # we have two connected bits
            $pdb_res   = $old_pdb_res+1;
            $pdb_chain = $first_chain;
        } else {
            $pdb_res   = $pdb_elem[5];
            $pdb_chain = $pdb_elem[4];
        }
        
        my $pdb_aa = $pdb_elem[3];
        
        if ($pdb_chain eq $first_chain) {
            
            $chain1{$pdb_res} = new residue;
            $chain1{$pdb_res}->chain     ($pdb_chain);
            $chain1{$pdb_res}->pdbresnum ($pdb_res);
            $chain1{$pdb_res}->aa        ($aa_3_to_1{lc $pdb_aa});
            
        } elsif ($pdb_chain eq $second_chain) {
            
            $chain2{$pdb_res} = new residue;
            $chain2{$pdb_res}->chain     ($pdb_chain);
            $chain2{$pdb_res}->pdbresnum ($pdb_res);
            $chain2{$pdb_res}->aa        ($aa_3_to_1{lc $pdb_aa});
            
        }
        
        $old_pdb_res = $pdb_res;
    }
    
    

    ##-----------------------------------------------------------------------##
    ## Processing the 2-chain Naccess file ----------------------------------##
    ##-----------------------------------------------------------------------##
    
    # produces "3SL9_A_C.rsa"
    unless (-e "pdb_partner_files/$pdb_code.rsa") {
        system($naccess . " pdb_partner_files/$pdb_code.pdb");
        system("mv $pdb_code.asa pdb_partner_files");
        system("mv $pdb_code.log pdb_partner_files");
        system("mv $pdb_code.rsa pdb_partner_files");
    }

    open(COMPLEXIN, "pdb_partner_files/$pdb_code.rsa") 
        or die ("Couldn't open the complex file from naccess");
    
    my %complexsasas;
    
    foreach (<COMPLEXIN>) {
        if (/^RES/) {
            my @line = split;
            # %complexsasas is a hash of hashes
            # first index: chain (e.g. "A")
            # second index: resid (e.g. 42)
            my $chn = $line[2];
            my $resid = $line[3];
            my $protocomplexsasas = $line[4];
            $complexsasas{$chn}{$resid} = $protocomplexsasas;
        }
    }
    
    close COMPLEXIN;
    
    ##-----------------------------------------------------------------------##
    ## Reading in the AlaScan output file -----------------------------------##
    ##-----------------------------------------------------------------------##
    
    ########################################################################
    # Switch this back later so it dies if no Rosetta output found, I guess?
    ########################################################################

    # Alanine scanning output in files like rosetta_output/1A0H_A_B_ala_out.txt
    my $ala_out = "rosetta_output/" . $pdb_code . "_ala_out.txt";
    my $okay_to_move = 0;
    unless (open(FILE, "<$ala_out")) {
        print STDERR "Can't open ala out file [$ala_out]. Skipping for now\n";
        next;
    }
    
    while (<FILE>) {
       
    
        # if no interface, get out!
        if (/centroids_by_jump_int called but no interface detected/) {
            # always move these
            next CODELOOP;
        }
        if (/overwrite/) {
            system("rm $ala_out");
            next CODELOOP;
        }
    
        # Look for lines that start AlaScan output
        unless (/protocols\.rosetta_scripts\.ParsedProtocol\.REPORT:\s+\w\s+\d+\s+\w+\s+\d+/) {
            next;
        }

        # Process this first, weird line
        
        # First, get rid of weirdness
        s/protocols\.rosetta_scripts\.ParsedProtocol\.REPORT:  //;
        
        # Break up the line into our desired fields
        my @lineobjs = split;
        
        # Index by residue number
        my $i = $lineobjs[1];
        my $j;
        
        # Hack to deal with incompletely written ala out files
        if (defined $chain1{$i}) {
            $chain1{$i}->rosresnum    ($lineobjs[3]);
            $chain1{$i}->ddg        ($lineobjs[4]);
            $chain1{$i}->dssp        ($lineobjs[5]);
        }
    
        # Process the rest of the (normal) lines in this block
        # but only if they match the last one's chain
        while (<FILE>) {
            
            # Skip if this is an empty line
            if (/^\s*$/) {
                next;
            }
            
            # Check whether this line matches the first one's chain
            
            @lineobjs = split;
            
            if ($lineobjs[0] eq $first_chain) {
                
                @lineobjs = split;
                
                # Index by residue number
                $i = $lineobjs[1];
                
                # Hack to deal with incompletely written ala out files
                if (defined $chain1{$i}) {
                    $chain1{$i}->rosresnum ($lineobjs[3]);
                    $chain1{$i}->ddg       ($lineobjs[4]);
                    $chain1{$i}->dssp      ($lineobjs[5]);
                } else {
                    next;
                }
                
                # ddg greater than 6 is unphysical for our purposes        
                # so set that to ">6" and don't count it for hotspots        
                # or toward averages...
                if ($lineobjs[4] ne 'n/a' && $lineobjs[4] > 6) {        
                    $chain1{$i}->ddg    (">6");        
                }
                
            } else {
                
                # Must be the second chain
                my @lineobjs2 = split;
                
                # Index by residue number
                if (defined $lineobjs2[1]) {
                    $j = $lineobjs2[1];
                
                    #Hack to deal with incompletely written ala out files
                    if (defined $chain2{$j}) {
                        $chain2{$j}->rosresnum ($lineobjs2[3]);
                        $chain2{$j}->ddg       ($lineobjs2[4]);
                        $chain2{$j}->dssp      ($lineobjs2[5]);
                    
                        # ddg greater than 6 is unphysical for our purposes        
                        # so set that to ">6" and don't count it for hotspots        
                        # or toward averages...        
                        if ($lineobjs2[4] ne 'n/a' && $lineobjs2[4] > 6) {        
                            $chain2{$j}->ddg    (">6");        
                        }
                    }
                }
            }
        }
        
        # Reached end of block.
        # ... and therefore end of what we care for in the alascan file
        last;
    }
    
    ##-----------------------------------------------------------------------##
    ## Initialize per-chain structures --------------------------------------##
    ##-----------------------------------------------------------------------##
    
    undef (%chains);
 
    $chains{$first_chain}  = \%chain1; # $chains{"A"} is reference to %chain1    
    my $chain1_len = scalar (keys %chain1);
    
    $chains{$second_chain} = \%chain2;
    my $chain2_len = scalar (keys %chain2);
    
    # initialize hotspots tracker to zero
    # takes nonzero value if it's a hotspot - indexed by chain and id/number
    undef %hotspots;
    
    foreach my $ch (keys %chains) {
        foreach my $id (keys %{$chains{$ch}}) {
            $hotspots{$ch}{$id} = 0;
        }
    }
    
    ##-----------------------------------------------------------------------##
    ## Per-chain processing -------------------------------------------------##
    ##-----------------------------------------------------------------------##
    
    # Recall: %chains was created
    # $chains{"A"} = \%chain1;
    
    foreach my $chn (sort { $a cmp $b } keys %chains) {
        
        my $pdb4lett = substr($pdb_code, 0, 4);
        
        # Perform Naccess's dSASA calculations on this file unless done already
        unless (-e "pdb_partner_files/$pdb4lett\_$chn.rsa") {
            if (-e "pdb_partner_files/$pdb4lett\_$chn.pdb") {
            system($naccess . " pdb_partner_files/$pdb4lett\_$chn.pdb");
            system ("mv $pdb4lett\_$chn.asa pdb_partner_files");
            system ("mv $pdb4lett\_$chn.log pdb_partner_files");
            system ("mv $pdb4lett\_$chn.rsa pdb_partner_files");
            } else {
                print STDERR "had to move on :-( \n";
                next;
            }
        }        
        
        # Process the result.
        open(CHAININ, "<pdb_partner_files/$pdb4lett\_$chn.rsa") or die;
        foreach (<CHAININ>) {
            if (/^RES/) {
                # this is a residue line
                my @line = split;
                
                my $resid = $line[3];
                my $sasa = $line[4];
                
                if (!exists $chains{$chn}{$resid}) {
                    next;
                }
                
                if (!defined $chains{$chn}{$resid}) {
                    next;
                }
                
                $chains{$chn}{$resid}->dsasa ($complexsasas{$chn}{$resid} - $sasa);
            }
        }
        close CHAININ;

        ##-------------------------------------------------------------------##
        ## Secondary structure and hotspot analysis--------------------------##
        ##-------------------------------------------------------------------##
        
        my $num_H_res = 0; # num of residues in a one-turn-plus helix
        my $secstruct = ""; # string to store the secondary structures present
        my $first_helix_pos = 0;
        my $last_helix_pos = 0;
    
        # first, determine what secondary structure there is at the interface
        # count the number of helical residues involved in (at least) a four-set of Hs
        
        #my @this_chains_keys = sort { $a <=> $b } keys %{$chains{$chn}};
        
        # Save keys into an array for easy access later. Remove residue IDs with letters
        # tacked on the end for now.
        undef(@this_chains_keys);
        foreach (keys %{$chains{$chn}}) {
            if ($_ !~ /[A-z]/) {
                push @this_chains_keys, $_;
            }
        }
        @this_chains_keys = sort { $a <=> $b } @this_chains_keys;
        
        
        for (my $n_chain_index = 1; $n_chain_index < (scalar @this_chains_keys) - 3; $n_chain_index++) {

            my $chain_i = $this_chains_keys[$n_chain_index];
            # Bit below is unnecessary now, because we remove ones with extra letters when
            # we create, the array, @this_chains_keys.
            # # Cut off extra letters in our residue. So 264A is treated just as 264.
            # $chain_i =~ s/[A-z]//g;
            
            my @H_bool;
            
            for (my $helix_c = $chain_i-3; $helix_c <= $chain_i+3; $helix_c++) {
                
                # Enter 0 if c doesn't exist in $chains{$chn}
                if (!exists ($chains{$chn}{$helix_c}) || !defined ($chains{$chn}{$helix_c}->dssp)) {
                    push (@H_bool, 0);
                } else {
                    if ($chains{$chn}{$helix_c}->dssp eq 'H') {
                        push (@H_bool, 1);
                    } else {
                        push (@H_bool, 0);
                    }
                }
            }
            
            # tracks if this residue is helical
            my $Hthis = 0;
            
            # Check i through i+3, elements 3 through 6 in @H_bool
            if (sum(@H_bool[3..6]) == 4) {
                if ($last_helix_pos == 0) {
                    $first_helix_pos = $chain_i;
                }
                $last_helix_pos = $chain_i; #tracking
                $Hthis = 1;
            } elsif (sum(@H_bool[2..5]) == 4 ||
                   sum(@H_bool[1..4]) == 4 || sum(@H_bool[0..3]) == 4) {
                $Hthis = 1;
                $last_helix_pos = $chain_i;
            }
                        
            $num_H_res += $Hthis; # it can be either one or zero
            
            # Record this as a hotspot if it meets the criterion
            # we separately check helicity so it's fine to do it blanket
            if (defined ($chains{$chn}{$chain_i}->ddg) &&
                         $chains{$chn}{$chain_i}->ddg ne 'n/a' &&
                         $chains{$chn}{$chain_i}->ddg ne '>6' &&
                         $chains{$chn}{$chain_i}->ddg >= 1) {
            
                $hotspots{$chn}{$chain_i} = 1;
            }
                
            # If THIS residue is not helical, then the last helix is dead
            if (!defined ($chains{$chn}{$chain_i}->ddg)
                       || $chains{$chn}{$chain_i}->dssp ne "H" 
                       || $chains{$chn}{$chain_i}->ddg eq 'n/a') {    #Decided against this
                
                if ($last_helix_pos > $first_helix_pos && $num_H_res >= 4) {
                    put_secstruct_to_DB ("alpha helix", $chn, $first_helix_pos, $last_helix_pos);
                }
                $num_H_res = 0;
                $last_helix_pos = 0;            
            }
        }
    } # end foreach chain
        
    close(FILE);
    if ($okay_to_move == 1) {    
        $okay_to_move = 0;
        print STDERR "Copied successfully completed $ala_out \n";
        system ("cp $ala_out ./rosetta_output_final/$pdb_code" . "_ala_out.txt");
    }
}

close (LIST);
exit 0;



sub put_secstruct_to_DB {
    
    my $secstruct = $_[0];
    my $chn = $_[1];
    my $first = $_[2];
    my $last = $_[3];
    
    my $total_chain_ddg = 0;            # total ddg
    my $total_chain_dsasa = 0;        # total dsasa
    my $total_hs_ddg = 0;        # total hotspot ddg, formerly $hsddg
    my $total_hs_dsasa = 0;    # total hotspot dsasa, formerly $hsdsasa
    my $total_SS_ddg = 0;
    my $total_SS_dsasa = 0;
        
    my $total_chain_hs = 0; # number hotspots
    my $total_SS_hs = 0;
        
    my $last_hs_pos=0;
    my $first_hs_pos=0;
    my $abs_SS_hs_pos_list = "";
    my $rel_SS_hs_pos_list = "i; ";
    
    my $arorascore;
    
    
    
    for (my $n_chain_index = 1; $n_chain_index < scalar @this_chains_keys; $n_chain_index++) {
    
        my $chain_i = $this_chains_keys[$n_chain_index];
            
        if (defined ($chains{$chn}{$chain_i}->ddg) && $chains{$chn}{$chain_i}->ddg =~ /^[\d\.\-]+$/) {
            $total_chain_ddg   += $chains{$chn}{$chain_i}->ddg;
        }
        if (defined ($chains{$chn}{$chain_i}->dsasa) && $chains{$chn}{$chain_i}->dsasa =~ /^[\d\.\-]+$/) {
            $total_chain_dsasa += $chains{$chn}{$chain_i}->dsasa;        
        }
            
        # special counters for helices        
        if ($secstruct eq "alpha helix") {    
            # not currently active because we actually only care about the longest helix
            # not the general collection of helical residues
        }
        
        if ($chain_i >= $first && $chain_i <= $last) {
            if (defined ($chains{$chn}{$chain_i}->ddg) && $chains{$chn}{$chain_i}->ddg =~ /^[\d\.\-]+$/) {    
                $total_SS_ddg += $chains{$chn}{$chain_i}->ddg;
            }
            if (defined ($chains{$chn}{$chain_i}->dsasa) && $chains{$chn}{$chain_i}->dsasa =~ /^[\d\.\-]+$/) {
                $total_SS_dsasa += $chains{$chn}{$chain_i}->dsasa;
            }
        }
            
            
        if ($hotspots{$chn}{$chain_i} == 1) {
            $total_chain_hs++;
            
            # build up relative HS location string
            #my $diff = $chain_i - $first_hs_pos;
            #if ($diff > 0) {
            #    $rel_SS_hs_pos_list .= "i + $diff; ";
            #}
            # special counters for SS hotspots
            
            if ($chain_i >= $first && $chain_i <= $last) {
            
                $total_hs_ddg   += $chains{$chn}{$chain_i}->ddg;
                if (defined ($chains{$chn}{$chain_i}->dsasa) && $chains{$chn}{$chain_i}->dsasa =~ /^[\d.\-]+$/) {
                    $total_hs_dsasa += $chains{$chn}{$chain_i}->dsasa;
                }
                $total_SS_hs++;
                
                $last_hs_pos = $chain_i;
                if ($first_hs_pos == 0) {
                    $first_hs_pos = $chain_i; # save the location of the first hs
                }
                        
                $abs_SS_hs_pos_list .= $chains{$chn}{$chain_i}->aa;
                $abs_SS_hs_pos_list .= $chains{$chn}{$chain_i}->pdbresnum;
                $abs_SS_hs_pos_list .= ", ";
                $abs_SS_hs_pos_list .= $chains{$chn}{$chain_i}->ddg;
                $abs_SS_hs_pos_list .= "; ";
                        
                my $diff = $chain_i - $first_hs_pos;
                if ($diff > 0) {
                    $rel_SS_hs_pos_list .= "i + $diff; ";
                }
            }
                    
        }
                
    }

        
    ##-------------------------------------------------------------------##
    ## Add to database --------------------------------------------------##
    ##-------------------------------------------------------------------##
        
    my $tablename = "helidb_tbl";
    
    # Some calculations for fields
        
    my $pdbc = substr $pdb_code, 0, 4;
        
    # The orientation of the two IDs in interfacechains field of the database
    #  changes. The first should be this chain and the second is its partner.
    my $interfacechains;
    if ($chn eq $first_chain) {
        $interfacechains =  substr ($pdb_code, 5, 1) . " " . substr ($pdb_code, 7, 1);
    } else {
        $interfacechains =  substr ($pdb_code, 7, 1) . " " . substr ($pdb_code, 5, 1);
    }
        
    #$chn; # just reminding us that it's here--no need to process further
        
    my $avg_SS_ddg;
    my $length = $last-$first+1;
    if ($length > 0) {
        $avg_SS_ddg = $total_SS_ddg / $length;
    } else {
        return;
    }
        
    my $pct_ddg_SS;
    if ($total_chain_ddg != 0) {
        $pct_ddg_SS = $total_SS_ddg / $total_chain_ddg;
    } else {
        $pct_ddg_SS = 'n/a';
    }
        
    my $avg_SS_dsasa;
        $avg_SS_dsasa = $total_SS_dsasa / $length;
        
    my $pct_dsasa_SS;
    if ($total_chain_dsasa != 0) {
        $pct_dsasa_SS = $total_SS_dsasa / $total_chain_dsasa;
    } else {
        $pct_dsasa_SS = 'n/a';
    }
    
    
    # Remove terminating ugly semicolons
    if (substr $abs_SS_hs_pos_list, -2 eq "; ") {        
        $abs_SS_hs_pos_list = substr $abs_SS_hs_pos_list, 0, -2;        
    }        
    if (substr $rel_SS_hs_pos_list, -2 eq "; ") {        
        $rel_SS_hs_pos_list = substr $rel_SS_hs_pos_list, 0, -2;        
    }
   
    my $SS_hs_distance=0; 
    if ($total_SS_hs > 1) {
        $SS_hs_distance = (split ' ', $rel_SS_hs_pos_list)[-1] + 1;
    }    
    
    my $SS_sequence;
    
    if ($first == 0 || $last == 0) {
            $SS_sequence = "n/a";
    } else {
        for (my $h = $first; $h <= $last; $h++) {
            if (defined ($chains{$chn}{$h})) {
                $SS_sequence .= $chains{$chn}{$h}->aa;
            }
        }
    }
        
    # more hotspots is better, scaling linearly        
    # more helix ddg is better, scaling linearly
    if ($avg_SS_ddg ne "n/a") {
        $arorascore = $total_SS_hs * $avg_SS_ddg;
    }
        
    # Round DDG averages, percent DDG helix, Arora score, percent helix dsasa, and        
    # average helix dsasa to 4 decimal places        
    $avg_SS_ddg   = sprintf ("%.4f", $avg_SS_ddg)   if ($avg_SS_ddg =~ /\d/);        
    $pct_ddg_SS   = sprintf ("%.4f", $pct_ddg_SS)   if ($pct_ddg_SS =~ /\d/);        
    $pct_dsasa_SS = sprintf ("%.4f", $pct_dsasa_SS) if ($pct_dsasa_SS =~ /\d/);        
    $avg_SS_dsasa = sprintf ("%.4f", $avg_SS_dsasa) if ($avg_SS_dsasa =~ /\d/);        
    $arorascore   = sprintf ("%.4f", $arorascore)   if ($arorascore =~ /\d/);
    
    
    my $SQL = "INSERT INTO $tablename (pdbcode, interfacechains, secstruct, ";
    $SQL   .=  "avg_SS_ddg, total_SS_ddg, total_chain_ddg, percent_SS_ddg, ";
    $SQL   .=  "avg_SS_dsasa, total_SS_dsasa, total_chain_dsasa, percent_SS_dsasa, ";
    $SQL   .=  "num_SS_hs, hs_SS_ids, hs_SS_relpos, hs_SS_len, ";
    $SQL   .=  "max_helix_length, first_SS_pos, last_SS_pos, SS_sequence, arorascore) ";
    
    $SQL   .= "VALUES (\"$pdbc\", \"$interfacechains\", \"$secstruct\", \"$avg_SS_ddg\", ";
    $SQL   .=  "\"$total_SS_ddg\", \"$total_chain_ddg\", \"$pct_ddg_SS\", ";
    $SQL   .=  "\"$avg_SS_dsasa\", \"$total_SS_dsasa\", \"$total_chain_dsasa\", ";
    $SQL   .=  "\"$pct_dsasa_SS\", \"$total_SS_hs\", \"$abs_SS_hs_pos_list\", ";
    $SQL   .=  "\"$rel_SS_hs_pos_list\", \"$SS_hs_distance\", ";
    $SQL   .=  "\"$length\", \"$first\", \"$last\", \"$SS_sequence\", ";
    $SQL   .=  "\"$arorascore\");";
        
    #$InsertRecord = $dbh->do($SQL);
        
    # Only put on DB if helical hotspotS        
    if ($total_SS_hs >= 2 && $length >= 4) {        
        print SQLOUT "GOOD: ";
        print "GOOD: ";
    }
    print SQLOUT $SQL . "\n"; # we be debuggin
    print $SQL . "\n";
}

close(SQLOUT);
